import os
from openai import OpenAI
from typing import Optional, Tuple
import xml.etree.ElementTree as ET
from dotenv import load_dotenv
import json
from datetime import datetime
from colorama import init, Fore, Style
from halo import Halo

from tools import TOOLS, execute_tool_call, get_tool_call
from prompts import get_system_prompt

# 初始化colorama
init()

# 加载环境变量
load_dotenv()

# 检查必要的环境变量
required_env_vars = ["OPENAI_API_KEY", "OPENAI_BASE_URL", "MODEL"]
missing_vars = [var for var in required_env_vars if not os.getenv(var)]
if missing_vars:
    raise ValueError(f"缺少必要的环境变量: {', '.join(missing_vars)}")

# 配置OpenAI客户端
client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL")
)
MODEL = os.getenv("MODEL")

class ChatHistory:
    def __init__(self, file_path="chat_history.json"):
        self.file_path = file_path
        self.current_session = []
        self.load_history()

    def load_history(self):
        """加载历史记录"""
        if os.path.exists(self.file_path):
            try:
                with open(self.file_path, "r", encoding="utf-8") as f:
                    self.history = json.load(f)
            except:
                self.history = []
        else:
            self.history = []

    def save_history(self):
        """保存历史记录"""
        if self.current_session:
            session_data = {
                "timestamp": datetime.now().isoformat(),
                "messages": self.current_session
            }
            self.history.append(session_data)
            
            with open(self.file_path, "w", encoding="utf-8") as f:
                json.dump(self.history, f, ensure_ascii=False, indent=2)

    def add_message(self, role: str, content: str):
        """添加新消息"""
        self.current_session.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })

    def clear_current_session(self):
        """清理当前会话"""
        self.current_session = []

def print_colored(text: str, role: str = "system"):
    """打印彩色文本"""
    color = {
        "system": Fore.WHITE,
        "user": Fore.GREEN,
        "assistant": Fore.CYAN,
        "tool": Fore.YELLOW,
        "error": Fore.RED
    }.get(role, Fore.WHITE)
    
    print(f"{color}{text}{Style.RESET_ALL}", end="", flush=True)

def chat_with_tools():
    """主对话循环"""
    history = ChatHistory()
    
    try:
        messages = [{"role": "system", "content": get_system_prompt()}]
        print_colored("AI工具助手已启动（输入'退出'结束对话）\n\n可用工具:\n", "system")
        for name, tool in TOOLS.items():
            print_colored(f"{name}: {tool.description}\n", "system")
        
        while True:
            try:
                user_input = input("\n用户: ").strip()
                if not user_input:
                    continue
                    
                if user_input.lower() in ['exit', 'quit', '退出']:
                    break
                
                # 添加用户消息到历史记录和消息列表
                history.add_message("user", user_input)
                messages.append({"role": "user", "content": user_input})
                
                while True:  # 工具调用循环
                    spinner = Halo(text="AI思考中", spinner='dots')
                    try:
                        spinner.start()
                        response = client.chat.completions.create(
                            model=MODEL,
                            messages=messages,
                            temperature=0,
                            stream=True
                        )
                        
                        # 获取第一个响应前显示加载动画
                        first_response = True
                        full_response = ""
                        
                        for chunk in response:
                            if chunk.choices[0].delta.content:
                                # 第一次收到响应时停止加载动画
                                if first_response:
                                    spinner.stop()
                                    print_colored("\nAI: ", "assistant")
                                    first_response = False
                                
                                content = chunk.choices[0].delta.content
                                print_colored(content, "assistant")
                                full_response += content
                                
                        print_colored("\n", "assistant")
                        
                        # 如果没有收到任何响应，停止加载动画
                        if first_response:
                            spinner.stop()
                        
                        history.add_message("assistant", full_response)
                        messages.append({"role": "assistant", "content": full_response})
                        
                        # 检查工具调用
                        tool_call = get_tool_call(full_response)
                        if tool_call is None:  # 如果没有工具调用，结束当前循环
                            break
                            
                        tool_name, xml_str = tool_call
                        result = execute_tool_call(xml_str)
                        print_colored(f"\n工具执行结果: {result}\n", "tool")
                        history.add_message("system", f"工具执行结果：\n{result}")
                        messages.append({"role": "system", "content": f"工具执行结果：\n{result}"})
                        
                        # 如果是完成任务的工具调用，结束对话
                        if tool_name == "complete_task":
                            break
                        
                    except Exception as e:
                        if spinner.spinner_id:
                            spinner.fail("AI响应出错")
                        error_msg = f"错误：{str(e)}"
                        print_colored(error_msg, "error")
                        history.add_message("error", error_msg)
                        break
                
            except KeyboardInterrupt:
                print_colored("\n已退出对话\n", "system")
                break
            except Exception as e:
                print_colored(f"\n发生错误：{str(e)}\n", "error")
                history.add_message("error", f"错误：{str(e)}")
                continue
                
    except KeyboardInterrupt:
        print_colored("\n程序被用户中断\n", "system")
    except Exception as e:
        print_colored(f"\n程序出错：{str(e)}\n", "error")
    finally:
        history.save_history()
        print_colored("\n会话结束\n", "system")

if __name__ == "__main__":
    chat_with_tools()